from pytube import YouTube
from urllib.parse import urlparse, parse_qs, urlunparse


def clean_url(url):
    # 只保留基础 YouTube 视频链接
    parsed = urlparse(url)
    if "youtu.be" in parsed.netloc:
        # 短链情况
        video_id = parsed.path.lstrip("/")
        return f"https://www.youtube.com/watch?v={video_id}"
    elif "youtube.com" in parsed.netloc:
        # 去掉多余参数，只保留 v=
        qs = parse_qs(parsed.query)
        if "v" in qs:
            return f"https://www.youtube.com/watch?v={qs['v'][0]}"
    return url


def download_youtube_video(url, save_path="."):
    try:
        url = clean_url(url)  # 清理一下
        yt = YouTube(url)
        print(f"🎬 视频标题: {yt.title}")
        print(f"👤 作者: {yt.author}")
        print(f"⏳ 长度: {yt.length} 秒")

        # 获取最高分辨率流
        stream = yt.streams.get_highest_resolution()
        print(f"⬇️ 正在下载: {stream.resolution}")

        stream.download(output_path=save_path)
        print(f"✅ 下载完成，保存到: {save_path}")
    except Exception as e:
        print(f"❌ 出错: {e}")


if __name__ == "__main__":
    url = input("请输入 YouTube 视频链接: ").strip()
    download_youtube_video(url, save_path=".")
